# coding: UTF-8
import torch
import torch.nn as nn
from transformers import AutoModel, AutoTokenizer, optimization


class Config(object):

    """配置参数"""
    def __init__(self, dataset):
        self.model_name = 'ERNIE'
        self.train_path = dataset + '/train.txt'                                # 训练集
        self.val_path = dataset + '/val.txt'                                    # 验证集
        self.test_path = dataset + '/test.txt'                                  # 测试集
        self.class_list = [x.strip() for x in open(
            dataset + '/class.txt').readlines()]                                # 类别名单
        self.save_path = './saved_dict/'+ dataset + '/' + self.model_name + '.ckpt'        # 模型训练结果
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')   # 设备

        self.require_improvement = 1000                                 # 若超过1000batch效果还没提升，则提前结束训练
        self.num_classes = len(self.class_list)                         # 类别数
        self.num_epochs = 5                                             # epoch数
        self.batch_size = 32                                            # mini-batch大小
        self.learning_rate = 5e-5                                       # 学习率
        self.bert_path = "nghuyong/ernie-1.0"
        self.tokenizer = AutoTokenizer.from_pretrained(self.bert_path)
        # print(self.tokenizer)
        self.hidden_size = 768


class Model(nn.Module):

    def __init__(self, config):
        super(Model, self).__init__()
        self.config = config
        self.bert = AutoModel.from_pretrained(config.bert_path)
        for param in self.bert.parameters():
            param.requires_grad = True
        self.fc = nn.Linear(config.hidden_size, config.num_classes)

    def forward(self, x):
        batch = self.config.tokenizer(x, padding=True, max_length=512, truncation=True, return_tensors="pt")
        _, pooled = self.bert(batch['input_ids'].to(self.config.device), attention_mask=batch['attention_mask'].to(self.config.device))
        out = self.fc(pooled)
        return out
